#include "tcp_config.hh"
#include "tcp_expectation.hh"
#include "tcp_fsm_test_harness.hh"
#include "tcp_header.hh"
#include "tcp_segment.hh"
#include "test_err_if.hh"
#include "util.hh"

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <exception>
#include <iostream>
#include <stdexcept>
#include <string>
#include <tuple>
#include <vector>

using namespace std;
using State = TCPTestHarness::State;

static constexpr unsigned NREPS = 32;

bool wrapping_lt(WrappingInt32 a, WrappingInt32 b) { return static_cast<int32_t>(b.raw_value() - a.raw_value()) > 0; }
bool wrapping_gt(WrappingInt32 a, WrappingInt32 b) { return wrapping_lt(b, a); }

int main() {
  try {
    TCPConfig cfg{};
    cfg.recv_capacity = 65000;
    auto rd = get_random_generator();

    // non-overlapping out-of-order segments
    for (unsigned rep_no = 0; rep_no < NREPS; ++rep_no) {
      const WrappingInt32 rx_isn(rd());
      const WrappingInt32 tx_isn(rd());
      TCPTestHarness test_1 = TCPTestHarness::in_established(cfg, tx_isn, rx_isn);
      vector<tuple<size_t, size_t>> seq_size;
      size_t datalen = 0;
      while (datalen < cfg.recv_capacity) {
        const size_t size =
            min(1 + (static_cast<size_t>(rd()) % (TCPConfig::MAX_PAYLOAD_SIZE - 1)), cfg.recv_capacity - datalen);
        seq_size.emplace_back(datalen, size);
        datalen += size;
      }
      shuffle(seq_size.begin(), seq_size.end(), rd);

      string d(datalen, 0);
      generate(d.begin(), d.end(), [&] { return rd(); });

      WrappingInt32 min_expect_ackno = rx_isn + 1;
      WrappingInt32 max_expect_ackno = rx_isn + 1;
      for (auto [off, sz] : seq_size) {
        test_1.send_data(rx_isn + 1 + off, tx_isn + 1, d.cbegin() + off, d.cbegin() + off + sz);
        if (off == min_expect_ackno.raw_value()) {
          min_expect_ackno = min_expect_ackno + sz;
        }
        max_expect_ackno = max_expect_ackno + sz;

        TCPSegment seg = test_1.expect_seg(ExpectSegment{}.with_ack(true), "test 1 failed: no ACK for datagram");

        auto &seg_hdr = seg.header();

        test_err_if(wrapping_lt(seg_hdr.ackno, min_expect_ackno) || wrapping_gt(seg_hdr.ackno, max_expect_ackno),
                    "test 1 failed: no ack or out of expected range");
      }

      test_1.execute(Tick(1));
      test_1.execute(ExpectData{}.with_data(d), "test 1 failed: got back the wrong data");
    }

    // overlapping out-of-order segments
    for (unsigned rep_no = 0; rep_no < NREPS; ++rep_no) {
      WrappingInt32 rx_isn(rd());
      WrappingInt32 tx_isn(rd());
      TCPTestHarness test_2 = TCPTestHarness::in_established(cfg, tx_isn, rx_isn);

      vector<tuple<size_t, size_t>> seq_size;
      size_t datalen = 0;
      while (datalen < cfg.recv_capacity) {
        const size_t size =
            min(1 + (static_cast<size_t>(rd()) % (TCPConfig::MAX_PAYLOAD_SIZE - 1)), cfg.recv_capacity - datalen);
        const size_t rem = TCPConfig::MAX_PAYLOAD_SIZE - size;
        size_t offs;
        if (rem == 0) {
          offs = 0;
        } else if (rem == 1) {
          offs = min(size_t(1), datalen);
        } else {
          offs = min(min(datalen, rem), 1 + (static_cast<size_t>(rd()) % (rem - 1)));
        }
        test_err_if(size + offs > TCPConfig::MAX_PAYLOAD_SIZE, "test 2 internal error: bad payload size");
        seq_size.emplace_back(datalen - offs, size + offs);
        datalen += size;
      }
      test_err_if(datalen > cfg.recv_capacity, "test 2 internal error: bad offset sequence");
      shuffle(seq_size.begin(), seq_size.end(), rd);

      string d(datalen, 0);
      generate(d.begin(), d.end(), [&] { return rd(); });

      WrappingInt32 min_expect_ackno = rx_isn + 1;
      WrappingInt32 max_expect_ackno = rx_isn + 1;
      for (auto [off, sz] : seq_size) {
        test_2.send_data(rx_isn + 1 + off, tx_isn + 1, d.cbegin() + off, d.cbegin() + off + sz);
        if (off <= min_expect_ackno.raw_value() && off + sz > min_expect_ackno.raw_value()) {
          min_expect_ackno = WrappingInt32(sz + off);
        }
        max_expect_ackno = max_expect_ackno + sz;  // really loose because of overlap

        TCPSegment seg = test_2.expect_seg(ExpectSegment{}.with_ack(true), "test 2 failed: no ACK for datagram");
        auto &seg_hdr = seg.header();

        test_err_if(wrapping_lt(seg_hdr.ackno, min_expect_ackno) || wrapping_gt(seg_hdr.ackno, max_expect_ackno),
                    "test 2 failed: no ack or out of expected range");
      }

      test_2.execute(Tick(1));
      test_2.execute(ExpectData{}.with_data(d), "test 1 failed: got back the wrong data");
    }
  } catch (const exception &e) {
    cerr << e.what() << endl;
    return err_num;
  }

  return EXIT_SUCCESS;
}
